import rclpy
from rclpy.node import Node
from nav_msgs.msg import Odometry
from geometry_msgs.msg import PoseStamped, TwistStamped, TransformStamped
from tf2_ros import TransformBroadcaster
from rclpy.qos import QoSProfile, ReliabilityPolicy
from rclpy.duration import Duration

class OdomCombiner(Node):

    def __init__(self):
        super().__init__('odom_combiner')
        
        qos_profile = QoSProfile(
            depth=5,
            reliability=ReliabilityPolicy.BEST_EFFORT
        )        
        # 存储最新的位姿和速度数据
        self.latest_pose = None
        self.latest_twist = None
        
        # 订阅两个输入话题
        self.pose_sub = self.create_subscription(
            PoseStamped,
            'ap/pose/filtered',
            self.pose_callback,
            qos_profile)
        
        self.twist_sub = self.create_subscription(
            TwistStamped,
            'ap/twist/filtered',
            self.twist_callback,
            qos_profile)
        
        # 创建输出话题发布器
        self.odom_pub = self.create_publisher(Odometry, 'odom', 10)
        
        # 创建TF广播器
        self.tf_broadcaster = TransformBroadcaster(self)
        
        # 设置固定发布频率为20Hz (0.05秒周期)
        self.publish_period = 0.05  # 秒
        self.timer = self.create_timer(self.publish_period, self.timer_callback)
        
        # 记录上一次发布时间，用于生成平滑的时间戳
        self.last_publish_time = self.get_clock().now()
        
        self.get_logger().info("优化后的Odometry组合器已启动，等待数据...")

    def pose_callback(self, msg):
        """存储最新位姿数据"""
        self.latest_pose = msg.pose

    def twist_callback(self, msg):
        """存储最新速度数据"""
        self.latest_twist = msg.twist

    def timer_callback(self):
        """以固定频率发布组合后的Odometry数据和TF变换"""
        # 检查是否有足够的数据
        if self.latest_pose is None or self.latest_twist is None:
            self.get_logger().debug('等待位姿和速度数据...', throttle_duration_sec=5)
            return
        
        # 生成平滑递增的时间戳，避免跳变
        current_time = self.last_publish_time + Duration(seconds=self.publish_period)
        self.last_publish_time = current_time
        
        # 创建并发布Odometry消息
        odom_msg = Odometry()
        odom_msg.header.stamp = current_time.to_msg()
        odom_msg.header.frame_id = 'odom'
        odom_msg.child_frame_id = 'base_footprint'
        odom_msg.pose.pose = self.latest_pose
        odom_msg.twist.twist = self.latest_twist
        self.odom_pub.publish(odom_msg)
        
        # 创建并发布TF变换 (odom -> base_footprint)
        t = TransformStamped()
        t.header.stamp = current_time.to_msg()
        t.header.frame_id = 'odom'
        t.child_frame_id = 'base_footprint'
        
        # 从位姿数据中提取变换信息
        t.transform.translation.x = self.latest_pose.position.x
        t.transform.translation.y = self.latest_pose.position.y
        t.transform.translation.z = self.latest_pose.position.z
        t.transform.rotation = self.latest_pose.orientation
        
        # 广播坐标变换信息
        self.tf_broadcaster.sendTransform(t)
        
        self.get_logger().info('发布odom消息和TF变换 (20Hz)', throttle_duration_sec=5)

def main(args=None):
    rclpy.init(args=args)
    node = OdomCombiner()
    try:
        rclpy.spin(node)
    except KeyboardInterrupt:
        pass
    finally:
        node.destroy_node()
        rclpy.shutdown()

if __name__ == '__main__':
    main()
